# -*- coding: utf-8 -*-
# @Time    : 2021/7/7 14:46
# @Author  : LuoTianHang


# ####################test.py 说明##########################
# this script is for the single test of the model

from typing import Union

import cv2
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from PIL import Image
from torch.backends import cudnn
from torchvision import transforms

from config import parser_config
from model import MCNN2 as MCNN
from otherutils import green_print, blue_print


def denormalize(im: Union[np.ndarray, torch.Tensor], mean=0.5, std=0.5):
    return im * std + mean


# data_transform = transforms.Compose([
#     transforms.Grayscale(3),
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# ])


def get_image(image, temp):
    get_density_map(image, temp)
    # NCHW->NHWC
    output = (denormalize(image.permute((0, 2, 3, 1)).
                          detach().
                          to('cpu').
                          numpy()) * 255).astype('uint8')
    output = output[0]
    output.reshape((output.shape[0], output.shape[1]))
    output = cv2.resize(output, (temp.width, temp.height))

    output = Image.fromarray(output)
    return output


def get_density_map(image, temp):
    image = F.relu(image)
    image = image.data.cpu().numpy()

    image = 255 * (image) / (np.max(image) + 1e-5)
    image = image[0][0]
    import matplotlib.pyplot as plt
    plt.imshow(image, cmap=plt.cm.jet, interpolation="bilinear")
    plt.savefig("result.jpg")
    plt.show()
    # plt.imshow(temp, cmap=plt.cm.jet, interpolation="bilinear")
    # plt.savefig("result.jpg")
    # plt.show()


def get_gt_map(image):
    image = image

    image = 255 * (image) / (np.max(image))

    import matplotlib.pyplot as plt
    plt.imshow(image, cmap=plt.cm.jet, interpolation="bilinear")
    plt.savefig("gt.jpg")
    plt.show()


class Inference(object):
    def __init__(self):
        self.args = parser_config()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()

        self.device = torch.device('cuda' if use_cuda else "cpu")

        self.model = MCNN()

        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count())).to(self.device)
            cudnn.enabled = True
            cudnn.benchmark = True

        blue_print("load the weight from pretrained-weight file")
        model_dict = self.model.state_dict()
        pretrained_dict = torch.load(self.args.pretrained_weight, map_location=self.device)['model_state_dict']
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
        model_dict.update(pretrained_dict)
        self.model.load_state_dict(model_dict)
        green_print("Finished to load the weight")

        self.model.eval()

    @torch.no_grad()
    def predict(self, path):
        den = pd.read_csv(path.replace('jpg', 'csv').replace('img', 'den')).values

        den = den.astype(np.float32, copy=False)
        img = Image.open(path).convert('L')

        w = den.shape[0]
        h = den.shape[1]

        wd_1 = img.width // 4
        ht_1 = img.height // 4
        den = cv2.resize(den, (wd_1, ht_1))
        den = den * ((w * h) / (wd_1 * ht_1))

        gt_label = np.sum(den)
        print("gt:", gt_label)
        img = Image.open(path).convert('L')
        temp = img
        img = torch.from_numpy(np.array(img,dtype=np.float32)).unsqueeze(0)
        img = img.unsqueeze(0).to(self.device)
        output = self.model(img)
        img2 = get_image(output, temp)
        img3 = get_gt_map(den)
        print("crowds:", torch.sum((output.cpu())).numpy())
        img3 = Image.new("RGB", (img2.width * 2, img2.height), (0, 0, 0))
        img3.paste(temp, (0, 0))
        img3.paste(img2, (img2.width, 0))
        # 拼接图片
        img3.save("hstack1.jpg")
        temp.convert('RGBA')
        img2.convert('RGBA')
        # 混合图片
        img4 = Image.blend(temp, img2, 0.9)
        img4.save("blend1.jpg")

if __name__=="__main__":
    model = Inference()
    model.predict('G:/datasets/crowdcounting/ProcessedData/shanghaitech_part_A/test/img/1.jpg')
